#include <stdio.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/err.h>

#include "sm4_crypto.h"

/* Function: sm4CbcEncrypt
 * Description:  SM4 Encrypt in CBC mode
 * Parameters:
 * pucKey           - pointer to encryption key buffer (must be 16 bytes)
 * stKeyLen         - encryption key length
 * pucIv            - pointer to IV buffer (must be 16 bytes)
 * stIvLen          - IV length
 * pucPlaintext     - pointer to plaintext buffer
 * stPlaintextLen   - plaintext Length
 * pucCiphertext    - pointer to ciphertext buffer
 * pstCiphertextLen - pointer to ciphertext length
 * Returns: 0 success, -1 failure
 */
int sm4CbcEncrypt(const unsigned char *pucKey,
                  size_t stKeyLen,
                  const unsigned char *pucIv,
                  size_t stIvLen,
                  const unsigned char *pucPlaintext,
                  size_t stPlaintextLen,
                  unsigned char *pucCiphertext,
                  size_t *pstCiphertextLen)
{
    EVP_CIPHER_CTX *pCipherCtx = NULL;
    int iLen = 0; 
    int iFinalLen = 0;

    /* Check key pointer*/
    if (!pucKey) {
        fprintf(stderr, "sm4CbcEncrypt: key buffer is null.\n");
        return (-1);
    }

    /* Check key length */
    if (stKeyLen != SM4_KEY_SIZE) {
        fprintf(stderr, "sm4CbcEncrypt: key length must be %d, but %zu.\n", SM4_KEY_SIZE, stKeyLen);
        return (-1);
    }

    /* Check IV pointer */
    if (!pucIv) {
        fprintf(stderr, "sm4CbcEncrypt: IV buffer is null.\n");
        return (-1);
    }

    /* Check IV length */
    if (stIvLen != SM4_KEY_SIZE) {
        fprintf(stderr, "sm4CbcEncrypt: IV length must be %d, but %zu.\n", SM4_KEY_SIZE, stIvLen);
        return (-1);
    }

    /* Check plaintext pointer */
    if (!pucPlaintext) {
        fprintf(stderr, "sm4CbcEncrypt: plaintext buffer is null.\n");
        return (-1);
    }

    /* plaintext length can be 0 */

    /* Check ciphertext pointer */
    if (!pucCiphertext) {
        fprintf(stderr, "sm4CbcEncrypt: ciphertext buffer is null.\n");
        return (-1);
    }

    /* Check ciphertext length pointer */
    if (!pstCiphertextLen) {
        fprintf(stderr, "sm4CbcEncrypt: ciphertext length buffer is null.\n");
        return (-1);
    }

    /* Create cipher context */
    pCipherCtx = EVP_CIPHER_CTX_new();
    if (!pCipherCtx) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize encrypt with SM4 CBC */
    if (EVP_EncryptInit_ex(pCipherCtx, EVP_sm4_cbc(), NULL, pucKey, pucIv) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Enable padding */
    EVP_CIPHER_CTX_set_padding(pCipherCtx, 1);

    /* Process plaintext */
    if (EVP_EncryptUpdate(pCipherCtx, pucCiphertext, &iLen, pucPlaintext, (int)stPlaintextLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Final encrypt */
    if (EVP_EncryptFinal_ex(pCipherCtx, pucCiphertext + iLen, &iFinalLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Check ciphertext buffer length */
    if ((*pstCiphertextLen) < (iLen + iFinalLen)) {
        fprintf(stderr, "sm4CbcEncrypt: ciphertext buffer length %zu is too small.\n", (*pstCiphertextLen));
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Set ciphertext length */
    *pstCiphertextLen = iLen + iFinalLen;

    /* Free cipher context */
    EVP_CIPHER_CTX_free(pCipherCtx);
 
    return (0);
}

/* Function: sm4CbcDecrypt
 * Description:  SM4 Decrypt in CBC mode
 * Parameters:
 * pucKey          - pointer to decryption key buufer (must be 16 bytes)
 * stKeyLen        - decryption key length
 * pucIv           - pointer to IV buffer (must be 16 bytes)
 * stIvLen         - IV length
 * pucCiphertext   - pointer to ciphertext buufer
 * stCiphertextLen - ciphertext length
 * pucPlaintext    - pointer to plaintext buffer
 * pstPlaintextLen - pointer to plaintext length
 * Returns: 0 success, -1 failure
 */
int sm4CbcDecrypt(const unsigned char *pucKey,
                  size_t stKeyLen,
                  const unsigned char *pucIv,
                  size_t stIvLen,
                  const unsigned char *pucCiphertext,
                  size_t stCiphertextLen,
                  unsigned char *pucPlaintext,
                  size_t *pstPlaintextLen)
{
    EVP_CIPHER_CTX *pCipherCtx = NULL;
    int iLen = 0;
    int iFinalLen = 0;

    /* Check key pointer */
    if (!pucKey) {
        fprintf(stderr, "sm4CbcDecrypt: key buffer is null.\n");
        return (-1);
    }

    /* Check key length */
    if (stKeyLen != SM4_KEY_SIZE) {
        fprintf(stderr, "sm4CbcDecrypt: key length must be %d, but %zu.\n", SM4_KEY_SIZE, stKeyLen);
        return (-1);
    }

    /* Check IV pointer */
    if (!pucIv) {
        fprintf(stderr, "sm4CbcDecrypt: IV buffer is null.\n");
        return (-1);
    }

    /* Check IV length */
    if (stIvLen != SM4_KEY_SIZE) {
        fprintf(stderr, "sm4CbcDecrypt: IV length must be %d, but %zu.\n", SM4_KEY_SIZE, stIvLen);
        return (-1);
    }

    /* Check ciphertext pointer */
    if (!pucCiphertext) {
        fprintf(stderr, "sm4CbcDecrypt: ciphertext buffer is null.\n");
        return (-1);
    }

    /* Check ciphertext length */
    if (stCiphertextLen == 0) {
        fprintf(stderr, "sm4CbcDecrypt: invalid ciphertext length %zu.\n", stCiphertextLen);
        return (-1);
    }

    /* Check plaintext pointer */
    if (!pucPlaintext) {
        fprintf(stderr, "sm4CbcDecrypt: plaintext buffer is null.\n");
        return (-1);
    }

    /* Check plaintext length pointer */
    if (!pstPlaintextLen) {
        fprintf(stderr, "sm4CbcDecrypt: plaintext length buffer is null.\n");
        return (-1);
    }

    /* Create cipher context */
    pCipherCtx = EVP_CIPHER_CTX_new();
    if (!pCipherCtx) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize decrypt with SM4 CBC */
    if (EVP_DecryptInit_ex(pCipherCtx, EVP_sm4_cbc(), NULL, pucKey, pucIv) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Enable padding */
    EVP_CIPHER_CTX_set_padding(pCipherCtx, 1);

    /* Process ciphertext */
    if (EVP_DecryptUpdate(pCipherCtx, pucPlaintext, &iLen, pucCiphertext, (int)stCiphertextLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Finalize decrypt */
    if (EVP_DecryptFinal_ex(pCipherCtx, pucPlaintext + iLen, &iFinalLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Check plaintext buffer length */
    if ((*pstPlaintextLen) < (iLen + iFinalLen)) {
        fprintf(stderr, "sm4CbcDncrypt: plaintext buffer length %zu is too small.\n", (*pstPlaintextLen));
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Set plaintext length */
    *pstPlaintextLen = iLen + iFinalLen;

    /* Clean cipher context */
    EVP_CIPHER_CTX_free(pCipherCtx);
    return (0);
}

/* Function: sm4EcbEncrypt
 * Description:  SM4 Encrypt in ECB mode
 * Parameters:
 * pucKey          - pointer to encryption key buffer (must be 16 bytes)
 * stKeyLen        - encryption key Length
 * pucPlaintext    - pointer to plaintext buffer
 * stPlaintextLen  - plaintext length
 * pucCiphertext   - pointer to ciphertext buffer
 * pstCiphertextLen- pointer to ciphertext length
 * Returns: 0 success, -1 failure
 */
int sm4EcbEncrypt(const unsigned char *pucKey,
                  size_t stKeyLen,
                  const unsigned char *pucPlaintext,
                  size_t stPlaintextLen,
                  unsigned char *pucCiphertext,
                  size_t *pstCiphertextLen)
{
    EVP_CIPHER_CTX *pCipherCtx = NULL;
    int iLen = 0;
    int iFinalLen = 0;

    /* Check key pointer */
    if (!pucKey) {
        fprintf(stderr, "sm4EcbEncrypt: key buffer is null.\n");
        return (-1);
    }

    /* Check key length */
    if (stKeyLen != SM4_KEY_SIZE) {
        fprintf(stderr, "sm4EcbEncrypt: key length must be %d, but %zu.\n", SM4_KEY_SIZE, stKeyLen);
        return (-1);
    }

    /* Check plaintext pointer */
    if (!pucPlaintext) {
        fprintf(stderr, "sm4EcbEncrypt: plaintext buffer is null.\n");
        return (-1);
    }

    /* plaintext length can be 0 */

    /* Check ciphertext pointer */
    if (!pucCiphertext) {
        fprintf(stderr, "sm4EcbEncrypt: ciphertext buffer is null.\n");
        return (-1);
    }

    /* Check ciphertext length pointer */
    if (!pstCiphertextLen) {
        fprintf(stderr, "sm4EcbEncrypt: ciphertext length buffer is null.\n");
        return (-1);
    }

    /* Create cipher context */
    pCipherCtx = EVP_CIPHER_CTX_new();
    if (!pCipherCtx) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize SM4-ECB encryption operation */
    if (EVP_EncryptInit_ex(pCipherCtx, EVP_sm4_ecb(), NULL, pucKey, NULL) != 1) {
        fprintf(stderr, "sm4EcbEncrypt: EVP_EncryptInit_ex failed");
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Enable padding */
    EVP_CIPHER_CTX_set_padding(pCipherCtx, 1);

    /* Process plaintext */
    if (EVP_EncryptUpdate(pCipherCtx, pucCiphertext, &iLen, pucPlaintext, (int)stPlaintextLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Finalize encrypt */
    if (EVP_EncryptFinal_ex(pCipherCtx, pucCiphertext + iLen, &iFinalLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Check ciphertext buffer length */
    if ((*pstCiphertextLen) < (iLen + iFinalLen)) {
        fprintf(stderr, "sm4EcbEncrypt: ciphertext buffer length %zu is too small.\n", (*pstCiphertextLen));
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Set ciphertext length */
    *pstCiphertextLen = iLen + iFinalLen;

    /* Clean context */
    EVP_CIPHER_CTX_free(pCipherCtx);
    return (0);
}

/* Function: sm4EcbDecrypt
 * Description: SM4 decrypt in ECB mode
 * Parameters:
 * pucKey          - pointer to decryption key buffer (must be 16 bytes)
 * stKeyLen        - decryption key length
 * pucCiphertext   - pointer to ciphertext buffer
 * stCiphertextLen - ciphertext length
 * pucPlaintext    - pointer to plaintext buffer
 * pstPlaintextLen - pointer to plaintext length
 * Returns: 0 success, -1 failure
 */
int sm4EcbDecrypt(const unsigned char *pucKey,
                  size_t stKeyLen,
                  const unsigned char *pucCiphertext,
                  size_t stCiphertextLen,
                  unsigned char *pucPlaintext,
                  size_t *pstPlaintextLen)
{
    EVP_CIPHER_CTX *pCipherCtx = NULL;
    int iLen = 0;
    int iFinalLen = 0;

    /* Check key pointer */
    if (!pucKey) {
        fprintf(stderr, "sm4EcbDecrypt: key buffer is null.\n");
        return (-1);
    }

    /* Check key length */
    if (stKeyLen != SM4_KEY_SIZE) {
        fprintf(stderr, "sm4EcbDecrypt: key length must be %d, but %zu.\n", SM4_KEY_SIZE, stKeyLen);
        return (-1);
    }

    /* Check ciphertext pointer */
    if (!pucCiphertext) {
        fprintf(stderr, "sm4EcbDecrypt: ciphertext buffer is null.\n");
        return (-1);
    }

    /* Check ciphertext length */
    if (stCiphertextLen == 0) {
        fprintf(stderr, "sm4EcbDecrypt: invalid ciphertext length %zu.\n", stCiphertextLen);
        return (-1);
    }

    /* Check plaintext pointer */
    if (!pucPlaintext) {
        fprintf(stderr, "sm4EcbDecrypt: plaintext buffer is null.\n");
        return (-1);
    }

    /* Check plaintext length pointer */
    if (!pstPlaintextLen) {
        fprintf(stderr, "sm4EcbDecrypt: plaintext length buffer is null.\n");
        return (-1);
    }

    /* Create cipher context */
    pCipherCtx = EVP_CIPHER_CTX_new();
    if (!pCipherCtx) {
        ERR_print_errors_fp(stderr);
        return (-1);
    }

    /* Initialize decrypt with SM4 ECB */
    if (EVP_DecryptInit_ex(pCipherCtx, EVP_sm4_ecb(), NULL, pucKey, NULL) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Enable padding */
    EVP_CIPHER_CTX_set_padding(pCipherCtx, 1);

    /* Process ciphertext */
    if (EVP_DecryptUpdate(pCipherCtx, pucPlaintext, &iLen, pucCiphertext, (int)stCiphertextLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Finalize decrypt */
    if (EVP_DecryptFinal_ex(pCipherCtx, pucPlaintext + iLen, &iFinalLen) != 1) {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Check plaintext buffer length */
    if ((*pstPlaintextLen) < (iLen + iFinalLen)) {
        fprintf(stderr, "sm4EcbDncrypt: plaintext buffer length %zu is too small.\n", (*pstPlaintextLen));
        EVP_CIPHER_CTX_free(pCipherCtx);
        return (-1);
    }

    /* Set plaintext length */
    *pstPlaintextLen = iLen + iFinalLen;

    /* Clean context */
    EVP_CIPHER_CTX_free(pCipherCtx);
    return (0);
}